package cn.spark.streaming;

import com.google.common.collect.ImmutableList;
import kafka.serializer.StringDecoder;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.api.java.*;
import org.apache.spark.streaming.kafka.KafkaUtils;
import org.elasticsearch.spark.rdd.api.java.JavaEsSpark;
import scala.Tuple2;

import java.util.*;

/**
 * Created by Niko Feng on 4/29/2017.
 */

/**
 * Spark Streaming for ElasticSearch Index Building
**/
public class ESSpark {
    final static SparkConf conf = new SparkConf()
            .setMaster("local[2]")
            .setAppName("ESSparkWordCount")
            .set("es.index.auto.create", "true")
            .set("es.nodes", "student28-x1")
            .set("es.port", "9200");

    final static JavaSparkContext jsc = new JavaSparkContext(conf);

    //Create a Spark Streaming Context
    final static JavaStreamingContext jssc = new JavaStreamingContext(jsc, Durations.seconds(5));

    public static void main(final String[] args) {
        try {
            //Create Kafka Configuration map
            Map<String, String> kafkaParams = new HashMap<String, String>();
            kafkaParams.put("metadata.broker.list",
                    "student39-x1:9092,student39-x2:9092,student40-x1:9092,student40-x2:9092");

            List<Set<String>> streamSets = new ArrayList<Set<String>>();

            Set<String> topics1 = new HashSet<String>();
            topics1.add(args[0].toString());
            streamSets.add(topics1);

            Set<String> topics2 = new HashSet<String>();
            topics2.add(args[1].toString());
            streamSets.add(topics2);

            Set<String> topics3 = new HashSet<String>();
            topics3.add(args[2].toString());
            streamSets.add(topics3);

            Set<String> topics4 = new HashSet<String>();
            topics4.add(args[3].toString());
            streamSets.add(topics4);

            //Create a list to store all the Kafka Streams for different topics
            int numStreams = 4;
            List<JavaPairDStream<String, String>> kafkaStreams = new ArrayList<JavaPairDStream<String, String>>(numStreams);
            for (int i = 0; i < numStreams; i++) {
                kafkaStreams.add(KafkaUtils.createDirectStream(
                        jssc,
                        String.class,
                        String.class,
                        StringDecoder.class,
                        StringDecoder.class,
                        kafkaParams,
                        streamSets.get(i)));
            }

            //Unify all the Kafka streams
            JavaPairDStream<String, String> unifiedStream = jssc.union(kafkaStreams.get(0), kafkaStreams.subList(1, kafkaStreams.size()));

            //Retreive the input raw data
            JavaDStream<String> wordsDstream = unifiedStream.flatMap(new FlatMapFunction<Tuple2<String, String>, String>() {
                @Override
                public Iterable<String> call(Tuple2<String, String> tuple) throws Exception {
                    return Arrays.asList(tuple._2());
                }
            });

            //Clean the raw data and transform the data into JSON format
            //Then save the JSON data to ElasticSearch
            JavaDStream<String> jsonDstream = wordsDstream.map(new Function<String, String>() {
                @Override
                public String call(String line) throws Exception {
                    System.out.println(line);
                    String[] tweets = line.split("\\|\\|");
                    String jsonStr = null;
                    if (tweets.length >= 2) {
                        jsonStr = ESSpark.transform2Json(tweets);
                        JavaRDD<String> stringRDD = jsc.parallelize(ImmutableList.of(jsonStr));
                        JavaEsSpark.saveJsonToEs(stringRDD, "tweets/obj");
                    }
                    return jsonStr;
                }
            });

            jsonDstream.print();

            jssc.start();
            jssc.awaitTermination();
            jssc.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            jsc.close();
            jssc.close();
        }
    }

    //Transform the raw data into JSON format
    public static String transform2Json(String[] tweets) {
        String jsonStr = "{" +
                "          \"text\": \"" + tweets[1] + "\"," +
//                "          \"text\": \"@MrKrimes Who will put Teknus in jail first? The Insurance Regulators or the SEC? $BLDV is a flat out P&amp;D scam. SEC TRADING HALT coming soon\",\n" +
//                "          \"is_quote_status\": false,\n" +
//                "          \"favorite_count\": 0,\n" +
//                "          \"retweeted\": false,\n" +
//                "          \"timestamp_ms\": \"1461695409657\",\n" +
//                "          \"entities\": {\n" +
//                "            \"entities_hashtags\": [],\n" +
//                "            \"entities_symbols\": [\n" +
//                "              {\n" +
//                "                \"entities_symbols_text\": \"BLDV\"\n" +
//                "              }\n" +
//                "            ],\n" +
//                "            \"entities_user_mentions\": [\n" +
//                "              {\n" +
//                "                \"entities_user_mentions_name\": \"MrKrimes\"\n" +
//                "              }\n" +
//                "            ]\n" +
//                "          },\n" +
//                "          \"retweet_count\": 0,\n" +
//                "          \"favorited\": false,\n" +
                "          \"user\": {" +
//                "            \"user_lang\": \"en\",\n" +
//                "            \"user_description\": null,\n" +
//                "            \"user_protected\": false,\n" +
//                "            \"user_friends_count\": 0,\n" +
//                "            \"user_verified\": false,\n" +
//                "            \"user_followers_count\": 0,\n" +
//                "            \"user_location\": null,\n" +
//                "            \"user_favourites_count\": 0,\n" +
//                "            \"user_name\": \"mark andrews\"," +
                "            \"user_name\": \"" + tweets[0] + "\"" +
//                "            \"user_created_at\": \"Tue Apr 26 14:39:47 +0000 2016\",\n" +
//                "            \"user_listed_count\": 0,\n" +
//                "            \"user_following\": null,\n" +
//                "            \"user_time_zone\": null,\n" +
//                "            \"user_statuses_count\": 5,\n" +
//                "            \"user_follow_request_sent\": null\n" +
                "          }," +
//                "          \"lang\": \"en\",\n" +
                "          \"created_at\": \"Tue Apr 26 18:30:09 +0000 2016\"" +
                "      }";
        return jsonStr;
    }
}
